import qualified Data.Map as Map

-- 0, 3, 6, 0, 3, 3, 1, 0, 4, 0, ...

type Index = Int
type Value = Int
type Memory = Map.Map Value [Index]
data State = State Value Memory deriving (Show)

updateMemory :: Memory -> (Index,Value) -> Memory
updateMemory memory (index, value) = 
    case indices of 
        Nothing -> Map.insert value [index] memory
        Just idxs -> Map.insert value (index:idxs) memory
    where
        indices = Map.lookup value memory

unwrap :: Maybe a -> a
unwrap (Just v) = v

step :: State -> Index -> State
step (State lastVal memory) idx 
    | length indices == 1 = State 0 (updateMemory memory (idx, 0))
    | otherwise           = State insertValue (updateMemory memory (idx, insertValue))
    where
        indices = unwrap $ Map.lookup lastVal memory

        lastIdx = head indices
        lastLastIdx = head . tail $ indices 
        insertValue = lastIdx - lastLastIdx

memoryGame :: [Value] -> State
memoryGame lst = foldl step initState [length lst..2019]
    where
        initMemory = foldl (\mem (i,v) -> Map.insert v [i] mem) Map.empty . zip [0..] $ lst
        initState = State (last lst) initMemory

getValue :: State -> Value
getValue (State value _) = value

main :: IO ()
main = do
    print . getValue . memoryGame $ [18,8,0,5,4,1,20]